import numpy as np

def predict_action(state: np.ndarray):
    """
    BLSW 策略（布林带 + KDJ + 成交量 + ATR波动率）
    - 自动计算 bb_lower
    - 阈值宽松，便于回测触发交易
    输入:
        state: np.ndarray of shape (b, 32, 20)
    输出:
        action: np.ndarray of shape (b, 1)，取值为 [-1, 0, 1]
    """
    b, l, c = state.shape
    assert l == 32 and c == 20

    # 指标索引
    IDX = {
        "close": 4,
        "volume": 5,
        "kdj_k": 13,
        "atr": 14,
        "bb_upper": 15
    }

    close = state[:, -1, IDX["close"]]
    close_window = state[:, :, IDX["close"]]  # shape (b, l)
    volume = state[:, -1, IDX["volume"]]
    kdj_k = state[:, -1, IDX["kdj_k"]]
    kdj_k_prev = state[:, -2, IDX["kdj_k"]]
    atr = state[:, -1, IDX["atr"]]
    atr_prev = state[:, -2, IDX["atr"]]

    # 计算布林带下轨
    bb_upper = state[:, -1, IDX["bb_upper"]]
    ma = close_window.mean(axis=1)
    std = (bb_upper - ma) / 2
    bb_lower = ma - 2 * std

    # 多头信号（阈值宽松）
    long_cond = (
        (close < bb_lower * 1.02) &   # 原来严格 < bb_lower，放宽 2%
        (kdj_k < 35) & (kdj_k > kdj_k_prev) &  # KDJ 放宽阈值
        (atr > atr_prev * 0.95) &  # ATR 放宽阈值
        (volume > volume.mean() * 0.8)  # 成交量放宽
    )

    # 空头信号（阈值宽松）
    short_cond = (
        (close > bb_upper * 0.98) &  # 放宽上轨阈值
        (kdj_k > 65) & (kdj_k < kdj_k_prev) &
        (atr > atr_prev * 0.95) &
        (volume > volume.mean() * 0.8)
    )

    action = np.zeros(b, dtype=np.int32)
    action[long_cond] = 1
    action[short_cond] = -1

    return action.reshape(-1, 1)
